from brian import *
from brian.library.random_processes import *
from brian.library.synapses import  *
import SuperNeuron as SuperNeuronCl
from MyConstants import *


mS = msiemens

# Conductances (mS/cm^2)
GNa = 120 * mS/cm**2
GK_dr = 100 * mS/cm**2
GCa_NS = 14 * mS/cm**2
GCa_ND = .3 * mS/cm**2
GK_CaS = 5 * mS/cm**2 # 3.136
GK_CaD = 1.1 * mS/cm**2  # 0.69
GCa_L = 0.33 * mS/cm**2
gleak = 0.51 * mS/cm**2
GNapD = 0.1 *mS/cm**2

# Static parameters
gc = 0.1  * mS/cm**2      # coupling conductance (mS/cm^2)
p = 0.1
Kd = 0.2 * 10**-2*mole*dmetre**-3     # uM
f = 0.01      # percent free to bound Ca
alpha = 0.009 *mole/(amp*second)*um**-1# mol/C/um
kca = 2 * ms**-1      # Ca removal rate
base_eqs=Equations('''
        dh/dt= (hinf-h)/Tauh : 1
        dn/dt= (ninf-n)/Taun : 1
        dmnS/dt = (mnSinf-mnS)/TaumN : 1
        dhnS/dt = (hnSinf-hnS)/TauhN : 1
        dmnD/dt = (mnDinf-mnD)/TaumN : 1
        dhnD/dt = (hnDinf-hnD)/TauhN : 1
        dmnap/dt = (mnapinf-mnap)/Taumnap : 1
        dhnap/dt = (hnapinf-hnap)/Tauhnap : 1
        dm/dt = (minf-m)/Taum : 1
        dml/dt = (mlinf-ml)/TaumL : 1
        dCaS/dt = f*(-alpha*ICaS-kca*CaS) : mole*dmetre**-3
        dCaD/dt = f*(-alpha*ICaD-kca*CaD) : mole*dmetre**-3
        dvm/dt = 1/C*(-INaS-IKS-ICaS-IleakS-IcouplingS-ISynI) : mV
        dvd/dt = 1/C*(-INapD-IKD-ICaD-IleakD-IcouplingD-ISynE-Irand) : mV
        
        Tauh = 30/(exp((vm/mV+50)/15)+exp(-(vm/mV+50)/16))*ms : ms
        Taun = 7/(exp((vm/mV+40)/40)+exp(-(vm/mV+40)/50))*ms : ms
        Tauhnap = 1200/(cosh((vm/mV + 59)/16))*ms :ms
        minf = 1/(1+exp((vm-Vhm)/Sm)) : 1
        hinf = 1/(1+exp((vm-Vhh)/Sh)) : 1
        ninf = 1/(1+exp((vm-Vhn)/Sn)) : 1
        mnapinf = 1/(1+exp((vd-Vhmnap)/Smnap)) : 1
        hnapinf = 1/(1+exp((vd-Vhhnap)/Shnap)) : 1
        mnSinf = 1/(1+exp((vm-VhmN)/SmN)) : 1
        hnSinf = 1/(1+exp((vm-VhhN)/ShN)) : 1
        mnDinf = 1/(1+exp((vd-VhmN)/SmN)) : 1
        hnDinf = 1/(1+exp((vd-VhhN)/ShN)) : 1
        mlinf = 1/(1+exp((vd-VhmL)/SmL)) : 1
        
        
        
        INapD = dyn*GNapD*mnap*hnap*(vd-ENa) : mA*umetre**-2
        INaS = dyn*GNa*m**3*h*(vm-ENa) : mA*umetre**-2
        IKS = dyn*(GK_dr*n**4 + GK_CaS*CaS/(CaS+Kd))*(vm-EK) : mA*umetre**-2
        ICaS = dyn*GCa_NS*mnS**2*hnS* (vm-ECa) : mA*umetre**-2
        IleakS = gleak*(vm-EleakS) : mA*umetre**-2
        IcouplingS = gc/p*(vm-vd) : mA*umetre**-2
        IKD = dyn*(GK_CaD*CaD/(CaD+Kd))*(vd-EK) : mA*umetre**-2
        ICaD = dyn*(GCa_ND*mnD**2*hnD+GCa_L*ml)*(vd-ECa) : mA*umetre**-2
        IleakD = gleak*(vd-EleakD) : mA*umetre**-2
        IcouplingD = gc/(1-p)*(vd-vm) : mA*umetre**-2
        
        ISynI = ipsp*(vd-Ei) : amp*umetre**-2
        ISynE = epsp*(vd-Ee) : amp*umetre**-2
        
        EleakS : mV
        EleakD : mV
        
        dyn : 1
        ''')



class MnNeuronFactory (SuperNeuronCl.SuperNeuron):
    
        
        
    def init(self, size,synapse='exp',Evariable=1,rand=False):
        self.size = size
        self.eqs = base_eqs
        self.rand = rand
        self.Evariable = Evariable
        self.add_synapses(synapse,tau1=tausyn)
        self.add_rand_curr(rand)
        
        self.Group=NeuronGroup(self.size,model=self.eqs,
                      threshold=EmpiricalThreshold(threshold=0*mV),
                      implicit=True)        
        self.reinit()
       
        
        return self.Group  
    
    def reinit(self):
        MnNeuronFactory.reinitNG(self.Group,self.Evariable)
    
    @staticmethod
    def reinitNG(Group,var):
        size = len(Group)
        ivs= -60*mV
        ivd= -60*mV
        Group.vm= ivs
        Group.vd= ivd
        Group.dyn = 1
        Group.h = 1/(1+exp((ivs-Vhh)/Sh))
        Group.n = 1/(1+exp((ivs-Vhn)/Sn))
        Group.mnS = 1/(1+exp((ivs-VhmN)/SmN))
        Group.hnS = 1/(1+exp((ivs-VhhN)/ShN))
        Group.mnD = 1/(1+exp((ivd-VhmN)/SmN))
        Group.hnD = 1/(1+exp((ivd-VhhN)/ShN))
        Group.ml = 1/(1+exp((ivd-VhmL)/SmL))
        Group.m = 1/(1+exp((ivs-Vhm)/Sm))
        Group.mnap = 1/(1+exp((ivd-Vhmnap)/Smnap))
        Group.hnap = 1/(1+exp((ivd-Vhhnap)/Shnap))
        Group.CaS = 0
        Group.CaD = 0
        Group.EleakS = (-65.0+randn(size)*6.5*var)*mV
        Group.EleakD = (-65.0+randn(size)*3.25*var)*mV

